# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de

import os.path as osp
import argparse
import pickle

import numpy as np
import torch
import open3d as o3d

import smplx


def main(
    model_folder,
    corr_fname,
    ext="npz",
    hand_color=(0.3, 0.3, 0.6),
    gender="neutral",
    hand="right",
):

    with open(corr_fname, "rb") as f:
        idxs_data = pickle.load(f)
        if hand == "both":
            hand_idxs = np.concatenate(
                [idxs_data["left_hand"], idxs_data["right_hand"]]
            )
        else:
            hand_idxs = idxs_data[f"{hand}_hand"]

    model = smplx.create(model_folder, model_type="smplx", gender=gender, ext=ext)
    betas = torch.zeros([1, 10], dtype=torch.float32)
    expression = torch.zeros([1, 10], dtype=torch.float32)

    output = model(betas=betas, expression=expression, return_verts=True)
    vertices = output.vertices.detach().cpu().numpy().squeeze()
    joints = output.joints.detach().cpu().numpy().squeeze()

    print("Vertices shape =", vertices.shape)
    print("Joints shape =", joints.shape)

    mesh = o3d.geometry.TriangleMesh()
    mesh.vertices = o3d.utility.Vector3dVector(vertices)
    mesh.triangles = o3d.utility.Vector3iVector(model.faces)
    mesh.compute_vertex_normals()

    colors = np.ones_like(vertices) * [0.3, 0.3, 0.3]
    colors[hand_idxs] = hand_color

    mesh.vertex_colors = o3d.utility.Vector3dVector(colors)

    o3d.visualization.draw_geometries([mesh])


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="SMPL-X Demo")

    parser.add_argument(
        "--model-folder", required=True, type=str, help="The path to the model folder"
    )
    parser.add_argument(
        "--corr-fname",
        required=True,
        type=str,
        dest="corr_fname",
        help="Filename with the hand correspondences",
    )
    parser.add_argument(
        "--gender", type=str, default="neutral", help="The gender of the model"
    )
    parser.add_argument(
        "--ext", type=str, default="npz", help="Which extension to use for loading"
    )
    parser.add_argument(
        "--hand",
        default="right",
        choices=["right", "left", "both"],
        type=str,
        help="Which hand to plot",
    )
    parser.add_argument(
        "--hand-color",
        type=float,
        nargs=3,
        dest="hand_color",
        default=(0.3, 0.3, 0.6),
        help="Color for the hand vertices",
    )

    args = parser.parse_args()

    model_folder = osp.expanduser(osp.expandvars(args.model_folder))
    corr_fname = args.corr_fname
    gender = args.gender
    ext = args.ext
    hand = args.hand
    hand_color = args.hand_color

    main(
        model_folder,
        corr_fname,
        ext=ext,
        hand_color=hand_color,
        gender=gender,
        hand=hand,
    )
